import sys
from pathlib import Path,PurePosixPath
from psycopg2.pool import ThreadedConnectionPool
import subprocess
import shutil,time,os
import gzip,bz2,lzma,zstandard
import paramiko
from faker import Faker
from queue import Queue

faker = Faker()


def get_client(ip, password, port=22):
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy)
    try:
        client.connect(hostname=ip, port=port, username="root", password=password, timeout=100,banner_timeout=600)
    except (
            paramiko.ssh_exception.NoValidConnectionsError,
            paramiko.ssh_exception.AuthenticationException,
            paramiko.ssh_exception.SSHException,
            TypeError,
            AttributeError,
    ) as e:
        print(f"无法连接到远程机器:{ip}.\n原因： {e}")
    return client



class RISC_V_UEFI:

    def __init__(self, **kwargs):
        self.arch = 'RISC-V'        # 当前测试类负责的指令集架构
        self.platform = 'UBOOT'     # 当前测试类负责的系统启动引导平台
        self.suite = kwargs.get('testsuite')             # 当前测试类待测试的mugen测试套名称
        self.case = kwargs.get('testcase')              # 当前测试类待测试的mugen测试名称
        self.vcpu = 2
        self.database_table_name = kwargs.get('database_table_name')
        self.workdir_runtime = kwargs.get('workdir_runtime')
        self.id_queue:Queue = kwargs.get('id_queue')
        self.multi_machine_lock = kwargs.get('multi_machine_lock')
        self.pool:ThreadedConnectionPool = kwargs.get('pgsql_pool')



    @staticmethod
    def make_openEuler_image(**kwargs):
        default_workdir:Path = kwargs.get('default_workdir')
        mugen_dir:Path = kwargs.get('mugen_dir')
        VIRT_VARS_FILE:Path = kwargs.get('VIRT_VARS_FILE')
        VIRT_CODE_FILE:Path = kwargs.get('VIRT_CODE_FILE')
        DRIVE_NAME:Path = Path(kwargs.get('DRIVE_FILE'))
        DRIVE_TYPE:Path = kwargs.get('DRIVE_TYPE')
        compress_format:str = kwargs.get('compress_format')
        print(compress_format)
        print(DRIVE_NAME)
        print(default_workdir)

        if compress_format == 'gzip':
            with gzip.open(default_workdir / DRIVE_NAME,'rb') as fin,open(default_workdir / Path(DRIVE_NAME).with_suffix(''),'wb') as fout:
                shutil.copyfileobj(fin, fout,length=1024*1024*32)
        elif compress_format == 'bzip2':
            with bz2.open(default_workdir / DRIVE_NAME,'rb') as fin,open(default_workdir / Path(DRIVE_NAME).with_suffix(''),'wb') as fout:
                shutil.copyfileobj(fin, fout,length=1024*1024*32)
        elif compress_format == 'xz':
            with lzma.open(default_workdir / DRIVE_NAME,'rb') as fin,open(default_workdir / Path(DRIVE_NAME).with_suffix(''),'wb') as fout:
                shutil.copyfileobj(fin, fout,length=1024*1024*32)
        elif compress_format == 'zstd':
            with zstandard.open(default_workdir / DRIVE_NAME,'rb') as fin,open(default_workdir / Path(DRIVE_NAME).with_suffix(''),'wb') as fout:
                shutil.copyfileobj(fin, fout,length=1024*1024*32)
        else:
            print("未检测到压缩格式，按照无压缩处理...")

        # 转变为绝对路径
        VIRT_VARS_FILE = VIRT_VARS_FILE.expanduser().resolve(strict=True)
        VIRT_CODE_FILE = VIRT_CODE_FILE.expanduser().resolve(strict=True)

        # 对两个固件进行软链接
        Path(default_workdir / PurePosixPath(VIRT_VARS_FILE).name).symlink_to(VIRT_VARS_FILE)
        Path(default_workdir / PurePosixPath(VIRT_CODE_FILE).name).symlink_to(VIRT_CODE_FILE)

        try:
            QEMU = subprocess.Popen(args=f"""
                qemu-system-riscv64 \
                  -nographic -machine virt,pflash0=pflash0,pflash1=pflash1,acpi=off \
                  -smp 8 -m 4G \
                  -object memory-backend-ram,size=2G,id=ram1 \
                  -numa node,memdev=ram1 \
                  -object memory-backend-ram,size=2G,id=ram2 \
                  -numa node,memdev=ram2 \
                  -blockdev node-name=pflash0,driver=file,read-only=on,filename={VIRT_CODE_FILE} \
                  -blockdev node-name=pflash1,driver=file,filename={VIRT_VARS_FILE} \
                  -drive file={default_workdir / DRIVE_NAME.with_suffix('')},format={DRIVE_TYPE},id=hd0,if=none \
                  -object rng-random,filename=/dev/urandom,id=rng0 \
                  -device virtio-vga \
                  -device virtio-rng-device,rng=rng0 \
                  -device virtio-blk-device,drive=hd0 \
                  -device virtio-net-device,netdev=usernet,mac={faker.mac_address()} \
                  -netdev user,id=usernet,hostfwd=tcp:127.0.0.1:20000-:22 \
                  -device qemu-xhci -usb -device usb-kbd -device usb-tablet
            """,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                shell=True,start_new_session=True
            )
            print(f"QEMU's pid = {QEMU.pid}")
        except subprocess.CalledProcessError as e:
            print(e)
            sys.exit(1)

        # QEMU启动RISC-V镜像需要较长一段时间
        time.sleep(60)
        subprocess.run(
            args = "nc -vz 127.0.0.1 20000",
            shell=True,
            stdout=subprocess.PIPE,
        )
        time.sleep(60)
        client:paramiko.SSHClient = get_client('127.0.0.1','openEuler12#$',20000)

        # copy mugen到镜像内
        with paramiko.SFTPClient.from_transport(client.get_transport()) as sftp:
            sftp.put(mugen_dir,'/root/')

        # 安装必备的rpm包
        stdin,stdout,stderr = client.exec_command(
            'dnf install -y git htop python3 && '
            'cd mugen/ && chmod +x dep_install.sh mugen.sh && bash dep_install.sh'
        )
        if stdout.channel.recv_exit_status() != 0:
            print(f"虚拟机中执行mugen初始化环境失败！报错信息:{stderr.read().decode('utf-8')}")
        print(stdout.read().decode('utf-8'))
        time.sleep(5)
        # copy mugen到镜像内(sftp只能传输文件而不能是目录)
        Path('/root/.ssh/known_hosts',exists_ok=True)
        scp = subprocess.run(
            args = f"export SSHPASS='openEuler12#$' && ssh-keygen -R '[localhost]:20000' && "
                   f"sshpass -e scp -r -P 20000 -o StrictHostKeyChecking=accept-new {mugen_dir} root@localhost:/root/",
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        if scp.returncode != 0:
            print(f"传输mugen进虚拟机失败.报错信息:{scp.stderr.decode()}")
            sys.exit(1)

        # 安装必备的rpm包
        stdin,stdout,stderr = client.exec_command(
            'set -e;'
            'dnf install -y git htop python3 && rpm --rebuilddb && dnf update -y && '
            'cd mugen/ && chmod +x dep_install.sh mugen.sh && bash dep_install.sh'
        )
        if stdout.channel.recv_exit_status() != 0:
            print(f"虚拟机中执行mugen初始化环境失败！报错信息:{stderr.read().decode('utf-8')}")

        stdin,stdout,stderr = client.exec_command(
            "systemctl disable --now firewalld && systemctl enable --now sshd"
        )
        if stdout.channel.recv_exit_status() != 0:
            print(f"关闭firewalld防火墙失败或者自启动sshd失败.报错信息:{stderr.read().decode('utf-8')}")

        stdin,stdout,stderr = client.exec_command(
            f"""
                nmcli con add type ethernet con-name net-static ifname enp0s3 ip4 10.0.0.2/24 gw4 10.0.0.254 && 
                nmcli con up net-static && nmcli device status && 
                cd mugen/ && bash mugen.sh -c --ip 10.0.0.2 --password openEuler12#$
            """
        )
        if stdout.channel.recv_exit_status() != 0:
            print(f"tap网络设置错误,或mugen创建conf/env.json失败!报错信息:{stderr.read().decode('utf-8')}")

        with client.open_sftp() as sftp:
            with sftp.open('/root/mugen/conf/env.json','r') as env:
                print(f"env content:{env.read().decode('utf-8')}")


        stdin,stdout,stderr = client.exec_command(
            "systemctl enable --now sshd && poweroff"
        )

        stdout.channel.recv_exit_status()
        # 初始化mugen
        # stdin,stdout,stderr = client.exec_command(
        #
        # )